import { ChatMessage } from "../message/chatMessage";
import { ModelBackend } from "../model";

/**
 * 打断消息
 * @param inputMessage 输入消息
 * @param modelBackend 模型后端
 * @param maxBreak 最大打断次数
 * @param CountHandelCallback 打断消息的回调函数，用以在token计数前处理消息，仅影响计数
 * @returns 打断后的消息数组
 */
export async function breakMessages(
    inputMessage: ChatMessage,
    modelBackend: ModelBackend,
    maxBreak: number,
    CountHandelCallback?: (message: ChatMessage, index: number) => Promise<string> | string
): Promise<ChatMessage[]> {
    const tokenLimit = modelBackend.getTokenLimit();
    const messagesToProcess: ChatMessage[] = [inputMessage];
    const result: ChatMessage[] = [];

    let processCount = 0;
    while (messagesToProcess.length > 0) {
        const message = messagesToProcess.shift()!;
        const prompt = CountHandelCallback
            ? await CountHandelCallback(message, processCount)
            : message.content;
        let tokenCount = modelBackend.getTokenNumberFromMessage(prompt);
        if (tokenCount > tokenLimit) {
            const cutLength = Math.ceil(message.content.length / 2);
            const nextMessages = Array.from({ length: 2 }, (_, k) => {
                const start = cutLength * k;
                const end = Math.min(start + cutLength, message.content.length);
                return new ChatMessage({
                    ...message,
                    content: message.content.slice(start, end)
                });
            });
            messagesToProcess.push(...nextMessages);
        } else {
            result.push(message);
        }
        processCount++;
        if (processCount >= maxBreak) {
            break;
        }
    }
    return result;
}
